import pandas as pd
import os
from datetime import datetime
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np


# --- 数据加载与基础处理函数 (保持不变) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据，并动态计算Volume Delta。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path)
                required_cols = ['open_time', 'open', 'high', 'low', 'close', 'volume', 'taker_buy_volume']
                if not all(col in df.columns for col in required_cols):
                    print(f"错误: 文件 {file_path} 缺少必要的列。")
                    continue

                for col in ['volume', 'taker_buy_volume']:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

                df['volume_delta'] = (2 * df['taker_buy_volume']) - df['volume']
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_data(df: pd.DataFrame, timeframe: str) -> pd.DataFrame:
    """处理数据，计算VWAP和CVD。"""
    if df.empty: return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'volume_delta']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    try:
        timeframe_minutes = int(timeframe.replace('m', ''))
        bars_in_24h = int((24 * 60) / timeframe_minutes)
    except ValueError:
        bars_in_24h = 288  # Default for 5m
    df['vwap_24h'] = (df['volume'] * ((df['high'] + df['low'] + df['close']) / 3)).rolling(window=bars_in_24h).sum() / \
                     df['volume'].rolling(window=bars_in_24h).sum()
    df['cvd_24h'] = df['volume_delta'].rolling(window=bars_in_24h).sum()
    df.dropna(subset=['vwap_24h', 'cvd_24h'], inplace=True)
    print("数据处理完成。")
    return df


# --- 绘图函数 (已移除S/R绘图逻辑) ---
def plot_dynamic_chart(df: pd.DataFrame, chart_title: str):
    """绘制K线、VWAP、成交量Delta和CVD图表。"""
    if df.empty:
        print("没有数据可供绘图。")
        return

    fig = make_subplots(rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.03, row_heights=[0.8, 0.2],
                        specs=[[{"secondary_y": True}], [{"secondary_y": False}]])

    # --- 基础图表元素 ---
    # K线图
    fig.add_trace(
        go.Candlestick(x=df.index, open=df['open'], high=df['high'], low=df['low'], close=df['close'], name='OHLC'),
        row=1, col=1)
    # 24小时 VWAP
    fig.add_trace(
        go.Scatter(x=df.index, y=df['vwap_24h'], name='24h VWAP', line=dict(color='orange', width=1.5, dash='dash')),
        row=1, col=1)
    # 逐笔成交量Delta (柱状图)
    fig.add_trace(go.Bar(x=df.index, y=df['volume_delta'], name='Volume Delta',
                         marker_color=['green' if v >= 0 else 'red' for v in df['volume_delta']]), secondary_y=True,
                  row=1, col=1)
    # 24小时滚动CVD (面积图)
    fig.add_trace(go.Scatter(x=df.index, y=df['cvd_24h'], name='24h Rolling CVD', line=dict(color='purple', width=1.5),
                             fill='tozeroy'), row=2, col=1)

    # --- 图表布局与样式 ---
    start_str, end_str = df.index[0].strftime('%Y-%m-%d'), df.index[-1].strftime('%Y-%m-%d')
    full_title = f'{chart_title}: K-Line with VWAP, Volume Delta & CVD ({start_str} to {end_str})'

    fig.update_layout(
        title_text=full_title,
        xaxis_rangeslider_visible=False,
        showlegend=True,
        hovermode='x unified',
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        yaxis=dict(title="Price (USDT)"),
        yaxis2=dict(title="Volume Delta", showgrid=False),
        yaxis3=dict(title="24h Rolling CVD"),
        xaxis2=dict(showticklabels=True)
    )
    fig.update_xaxes(rangeslider_visible=False)
    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "ETHUSDT"
    TIMEFRAME = "5m"
    START_TIME = "2025-02-01 00:00:00"
    END_TIME = "2025-02-28 00:00:00"
    # ====================================================

    print("--- K线图表程序启动 ---")

    # 直接使用配置的时间范围加载数据
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)

    if not raw_data_df.empty:
        # 处理数据以计算指标
        processed_df = process_data(raw_data_df, TIMEFRAME)

        if not processed_df.empty:
            # 筛选出在显示范围内的数据进行绘图 (处理后数据可能会因rolling而减少)
            chart_df = processed_df.loc[START_TIME:END_TIME]
            chart_title_str = f"{SYMBOL} - {TIMEFRAME}"

            # 绘制图表
            plot_dynamic_chart(chart_df, chart_title=chart_title_str)

            print("--- 图表生成完毕 ---")
        else:
            print("\n数据处理后为空，无法生成图表。可能是因为数据量不足以计算滚动指标。")
    else:
        print("\n未能加载任何数据，请检查配置。")